import sys

import torch
import torch.nn as nn
from abc import ABC, abstractmethod

if sys.version_info >= (3, 10):
    from .functional import DWT1D, IDWT1D, DWT2D, IDWT2D, DWT3D, IDWT3D, _to_wavelet_coefs
else:
    from .functional_compatible import DWT1D, IDWT1D, DWT2D, IDWT2D, DWT3D, IDWT3D, _to_wavelet_coefs

"""
To Pytorch Layers
"""


class DWTBase(ABC, nn.Module):
    def __init__(self, wavelet_level: int = 1, wave='haar'):
        super().__init__()
        lohi = _to_wavelet_coefs(wave)
        self.register_buffer('lohi', lohi)
        self.wavelet_level = wavelet_level

    @abstractmethod
    def forward(self, x):
        pass


class DWTForward3D(DWTBase):
    def __init__(self, wavelet_level: int = 1, wave='haar'):
        super(DWTForward3D, self).__init__(wavelet_level, wave)

    def forward(self, x):
        yh = []
        ll = x
        for _ in range(self.wavelet_level):
            coefs = DWT3D.apply(ll, self.lohi)
            ll = coefs[:, 0, ]
            high = coefs[:, 1:, ]
            yh.append(high)
        return ll, yh


class DWTInverse3D(DWTBase):
    def __init__(self, wavelet_level: int = 1, wave='haar'):
        super(DWTInverse3D, self).__init__(wavelet_level, wave)

    def forward(self, coefs):
        yl, yh = coefs
        ll = yl
        for i in range(self.wavelet_level - 1, -1, -1):
            x = torch.cat([ll.unsqueeze(1), yh[i]], dim=1)
            ll = IDWT3D.apply(x, self.lohi)
        return ll


class DWTForward2D(DWTBase):
    def __init__(self, wavelet_level: int = 1, wave='haar'):
        super(DWTForward2D, self).__init__(wavelet_level, wave)

    def forward(self, x):
        yh = []
        ll = x
        for _ in range(self.wavelet_level):
            coefs = DWT2D.apply(ll, self.lohi)
            ll = coefs[:, 0, ]
            high = coefs[:, 1:, ]
            yh.append(high)
        return ll, yh


class DWTInverse2D(DWTBase):
    def __init__(self, wavelet_level: int = 1, wave='haar'):
        super(DWTInverse2D, self).__init__(wavelet_level, wave)

    def forward(self, coefs):
        yl, yh = coefs
        ll = yl
        for i in range(self.wavelet_level - 1, -1, -1):
            x = torch.cat([ll.unsqueeze(1), yh[i]], dim=1)
            ll = IDWT2D.apply(x, self.lohi)
        return ll


class DWTForward1D(DWTBase):
    def __init__(self, wavelet_level: int = 1, wave='haar'):
        super(DWTForward1D, self).__init__(wavelet_level, wave)

    def forward(self, x):
        yh = []
        ll = x
        for _ in range(self.wavelet_level):
            coefs = DWT1D.apply(ll, self.lohi)
            ll = coefs[:, 0, ]
            high = coefs[:, 1:, ]
            yh.append(high)
        return ll, yh


class DWTInverse1D(DWTBase):
    def __init__(self, wavelet_level: int = 1, wave='haar'):
        super(DWTInverse1D, self).__init__(wavelet_level, wave)

    def forward(self, coefs):
        yl, yh = coefs
        ll = yl
        for i in range(self.wavelet_level - 1, -1, -1):
            x = torch.cat([ll.unsqueeze(1), yh[i]], dim=1)
            ll = IDWT1D.apply(x, self.lohi)
        return ll
